#!/usr/bin/env python

# Entension of tutorial bot that gobbles food
#
# Author: Bennett Wineholt
#
# Collaborative Diffusion chosen from reading Benjamin Meyer
#     with modifications envisioned by Gautam Kamath

import sys
import pickle

from copy import deepcopy
from random import shuffle
from scipy import *
from scipy import ndimage

from ants import *

TRANSITION = 0.125
REMAIN = 0.5

INFLUENCE_SPREAD = 10
FOOD_INFLUENCE = (100/TRANSITION)**INFLUENCE_SPREAD


class MyBot:
    def __init__(self):
        pass
 
    def do_setup(self, ants):
        self.unseen = []
        for row in range(ants.rows):
            for col in range(ants.cols):
                self.unseen.append((row, col))
        # initialize one ant step filter
        self.one_step = array([[0,TRANSITION,0],
                               [TRANSITION,REMAIN,TRANSITION],
                               [0,TRANSITION,0]])
        # land pixels adjacent to pixels never seen
        self.frontier = None
        self.hills = [] 
        # total influence score of each map pixel
        self.influence_map = zeros((ants.rows, ants.cols))
    
    def do_turn(self, ants):

        # Begin helper method declarations

        def do_move_direction(loc, direction):
            new_loc = ants.destination(loc, direction)
            if (ants.unoccupied(new_loc) and new_loc not in orders):
                ants.issue_order((loc, direction))
                orders[new_loc] = loc
                targets[new_loc] = loc
                return True
            else:
                return False

        def do_move_location(loc, dest):
            directions = ants.direction(loc, dest)
            for direction in directions:
                if do_move_direction(loc, direction):
                    targets[dest] = loc
                    return True
            return False

        # Returns complement of known water on map
        def find_not_water():
            # matrix size of map marking pixels that are not water
            return ((array(ants.map) - WATER) != 0).astype(int)

        # propagates influence along land by one step
        #
        # inflation corrects for the loss of value given to neighbors since
        # some neighbors are water, 0 sink
        #
        # similarly, to utilize desirable properties of the normalized kernel
        # we re-adjust for the inflated self-influence and maintain balance
        #
        # Requires all scipy array
        #
        # Returns updated map
        def propagate_influence(imap):
            return ((ndimage.correlate(imap, self.one_step)) *
                    self.not_water)

        # update the influence map for a turn with values of each location
        # based on
        # - food
        def process_influence_map():
            imap = zeros((ants.rows, ants.cols))
            for (food_r, food_c) in ants.food():
                imap[food_r][food_c] = FOOD_INFLUENCE
            for i in range(0, INFLUENCE_SPREAD):
                imap = propagate_influence(imap)
            return imap

         # update influence map for food

        # End helper method declarations


        # orders to return for each ant
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        orders = {}
        # destination targets to prevent conflicting pursuit of food
        # dict of ((dest_r, dest_c), (ant_r, ant_c))
        targets = {}

        # update visible map locations
        for loc in self.unseen[:]:
            if ants.visible(loc):
                self.unseen.remove(loc)

        # update obstacle map, aka what's not water
        self.not_water = find_not_water()

        # influence map accumulates value of map locations
        self.influence_map = process_influence_map()

        # prevent stepping on own hill
        for hill_loc in ants.my_hills():
            orders[hill_loc] = None

        # find close food
        #ant_dist = []
        #food_list = ants.food()
        #shuffle(food_list)
        #for food_loc in food_list:
        #    for ant_loc in ants.my_ants():
        #        dist = ants.distance(ant_loc, food_loc)
        #        ant_dist.append((dist, ant_loc, food_loc))
        #ant_dist.sort()
        #for dist, ant_loc, food_loc in ant_dist:
        #    if food_loc not in targets and ant_loc not in targets.values():
        #        do_move_location(ant_loc, food_loc)

        # find food and explore based on collaborative diffusion
        for (ant_r, ant_c) in ants.my_ants():
            pos_dest = []
            for direction in list(AIM.keys()):
                (new_r, new_c) = ants.destination((ant_r, ant_c), direction)
                pos_dest.append((self.influence_map[new_r][new_c],
                                 (new_r, new_c), direction))
            list.sort(pos_dest, reverse=True)
            #if pos_dest[0][0] < NEGLIGIBLE_INFLUENCE:
            #    break
            for (score, dest, direction) in pos_dest:
                if (dest not in targets and 
                        (ant_r, ant_c) not in targets.values() and
                        ants.unoccupied(dest)):
                    if do_move_direction((ant_r, ant_c), direction):
                        break

        # attack hills
        for hill_loc, hill_owner in ants.enemy_hills():
            if hill_loc not in self.hills:
                self.hills.append(hill_loc)        
        ant_dist = []
        for hill_loc in self.hills:
            for ant_loc in ants.my_ants():
                if ant_loc not in orders.values():
                    dist = ants.distance(ant_loc, hill_loc)
                    ant_dist.append((dist, ant_loc, hill_loc))
        ant_dist.sort()
        for dist, ant_loc, hill_loc in ant_dist:
            do_move_location(ant_loc, hill_loc)
       
        # random move
        # has same effect as exploring unseen without worrying about pathing
        #
        # TODO(bmw): is 2 u-sec enough to terminate?
        for ant_loc in ants.my_ants():
            if ant_loc not in orders.values(): 
                if (ants.time_remaining() < 2):
                    break
                directions = list(AIM.keys())
                shuffle(directions)
                for direction in directions:
                    dest = ants.destination(ant_loc, direction)
                    if (not dest in targets and
                            ants.unoccupied(dest)):
                        if do_move_location(ant_loc, dest):
                            break
        
        # unblock own hill
        for hill_loc in ants.my_hills():
            if hill_loc in ants.my_ants() and hill_loc not in orders.values():
                for direction in ('s','e','w','n'):
                    if do_move_direction(hill_loc, direction):
                        break

        
if __name__ == '__main__':
    # psyco will speed up python a little, but is not needed
    try:
        import psyco
        psyco.full()
    except ImportError:
        pass
    
    try:
        # if run is passed a class with a do_turn method, it will do the work
        # this is not needed, in which case you will need to write your own
        # parsing function and your own game state class
        Ants.run(MyBot())
    except KeyboardInterrupt:
        print('ctrl-c, leaving ...')
